#!/bin/bash

set -e -o pipefail

BASE=$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")
FLASH_TOOL="$BASE/flash-tool"
KHADAS_TOOL="/usr/local/bin/$(basename $0)"

if [ "$(uname -s)" = "Darwin" ]; then
	ADNL_TOOL="$BASE/tools/adnl/macos/adnl_burn_pkg"
else
	ADNL_TOOL="$BASE/tools/adnl/adnl_burn_pkg"
fi

IMAGE=
DEVICE=
BOARD=
IMAGE_INSTALL_TYPE=

RED='\033[0;31m'
RESET='\033[m'

error_msg() {
	echo -e "$RED"ERROR:"$RESET" $1
}

usage() {
	echo -e "Usage:"
	echo -e "emmc:   $0 -b <VIM1|VIM2|VIM3|VIM4> -i <path-to-image> [-n]"
	echo -e "sdcard: $0 -d </dev/sdX> -i <path-to-image>"
	echo -e "\nOptions:"
	echo -e "  -n    Do not erase the whole flash chip (default to erase)."
	echo -e "  -s    Skip the USB connectivity check, it will be verified more reliably later during the flashing process (default to check)."
}

flash_sdcard() {
	pv -tpreb ${IMAGE} | sudo dd of=${DEVICE} conv=notrunc
	echo "Sync..."
	sync
}

## Calculate time
## $1 - time in seconds
time_cal() {
	local minutes

	minutes=$(($1 / 60))

	echo "Time elapsed: $minutes minute(s)."
}

if [ ! -L $KHADAS_TOOL ]; then
	error_msg "Please install `basename $0`. Execute 'INSTALL' script to install it."
	exit 1
fi

SKIP_USB_CHECK="no"

# Default to erase if the option is not explicitly provided
# to match the default behavior of $FLASH_TOOL and $ADNL_TOOL
ERASE_ALL="yes"

while getopts "d:i:b:Dhrns" flag; do
	case $flag in
		d)
		DEVICE="$OPTARG"
		;;
		i)
		IMAGE="$OPTARG"
		;;
		b)
		BOARD="$OPTARG"
		;;
		D)
		DEBUG="--debug"
		;;
		r)
		RESET_BOARD="yes"
		;;
		n)
		ERASE_ALL="no"
		;;
		s)
		SKIP_USB_CHECK="yes"
		;;
		h)
		usage
		exit
		;;
	esac
done

if [ ! -f "$IMAGE" ]; then
	error_msg "Image '$IMAGE' doesn't exist!"
	usage
	exit -1
fi

case $IMAGE in
	*.xz)
	echo "Decompressing image ..."
	xz -d -f $IMAGE
	IMAGE=${IMAGE%.*}
	;;
esac

if [ -z "$BOARD" ]; then
	BOARD="VIM1"
fi

if [ "$(uname -s)" = "Darwin" ]; then
	if fdisk -t "$IMAGE" && ! fdisk -d "$IMAGE" | grep -q 0xEE; then
		IMAGE_INSTALL_TYPE="SD-USB"
	else
		IMAGE_INSTALL_TYPE="EMMC"
	fi
else
	partition=$(fdisk -l "$IMAGE" | grep "Disklabel type" | awk -F ": " '{print $2}' || true)
	if [ "$partition" == "dos" ]; then
		IMAGE_INSTALL_TYPE="SD-USB"
	else
		IMAGE_INSTALL_TYPE="EMMC"
	fi
fi

start_time=`date +%s`

if [ $DEVICE ]; then
	if [ ! -b $DEVICE ]; then
		error_msg "'$DEVICE' is not a block device! Please make sure the device you choose is right."
		exit -1
	fi

	if [ "$IMAGE_INSTALL_TYPE" != "SD-USB" ]; then
		error_msg "Try to burn to SD/USB storage,but the image installation type is '$IMAGE_INSTALL_TYPE', please use 'SD-USB' image!"
		exit -1
	fi

	echo "Burning image '$IMAGE' to SD/USB storage..."
	flash_sdcard
else
	if [ "$IMAGE_INSTALL_TYPE" != "EMMC" ]; then
		error_msg "Try to burn to eMMC storage, but the image installation type is '$IMAGE_INSTALL_TYPE', please use 'EMMC' image!"
		exit -1
	fi

	if [ "$SKIP_USB_CHECK" != "yes" ]; then
		if [ "$(uname -s)" = "Darwin" ]; then
			if ! ioreg -p IOUSB -l | grep -q "Amlogic" > /dev/null; then
				error_msg "You should put your board enter upgrade mode!"
				exit -1
			fi
		else
			if ! lsusb | grep -q "Amlogic, Inc." > /dev/null; then
				error_msg "You should put your board enter upgrade mode!"
				exit -1
			fi
		fi
	fi

	if [ "$BOARD" == "VIM1" ] || [ "$BOARD" == "VIM2" ]; then
		SOC="gxl"
		DISPLAY_BOARD="VIM1/VIM2"
	elif [ "$BOARD" == "VIM3" ]; then
		SOC="g12a"
		DISPLAY_BOARD="VIM3/VIM3L"
	elif [ "$BOARD" == "VIM4" ]; then
		SOC="t7"
		DISPLAY_BOARD="VIM4"
	elif [ "$BOARD" == "VIM1S" ]; then
		SOC="s4"
		DISPLAY_BOARD="VIM1S"
	else
		error_msg "Unsupported board: [$BOARD]"
		exit -1
	fi

	echo "Burning image '$IMAGE' for '$DISPLAY_BOARD' to eMMC..."
	if [ "$BOARD" == "VIM1" ] || [ "$BOARD" == "VIM2" ] || [ "$BOARD" == "VIM3" ]; then
		if [ "$RESET_BOARD" == "yes" ]; then
			RESET_BOARD="--reset=y"
		else
			RESET_BOARD=
		fi
		if [ "$ERASE_ALL" == "yes" ]; then
			WIPE_OPTION="--wipe"
		else
			WIPE_OPTION=
		fi

		$FLASH_TOOL --img=$IMAGE --parts=all ${WIPE_OPTION} --soc=$SOC ${RESET_BOARD} ${DEBUG}
	elif [ "$BOARD" == "VIM4" ] || [ "$BOARD" == "VIM1S" ]; then
		if [ "$RESET_BOARD" == "yes" ]; then
			RESET_BOARD="-r 1"
		else
			RESET_BOARD=
		fi
		if [ "$ERASE_ALL" == "yes" ]; then
			ERASE_FLASH="-e 1"
		else
			ERASE_FLASH="-e 0"
		fi
		$ADNL_TOOL -p $IMAGE ${RESET_BOARD} ${ERASE_FLASH}
	fi
fi

end_time=`date +%s`

time_cal $(($end_time - $start_time))

echo "Done!"
date
